# -*- coding: utf-8 -*-
"""
Created on Thu Apr 15 10:49:24 2021

@author: Lenovo-pc
"""

from xgboost.sklearn import XGBRegressor
from sklearn.model_selection import GridSearchCV
import joblib
import numpy as np

import warnings
warnings.filterwarnings("ignore")

class xgb_gbdt:
    def __init__(self):
        pass
    
    def train(self,train_data, train_labels):
        # param_grid = {
        #     'max_depth':range(10,100,10)
        #     }

        # gsearch = GridSearchCV(
        #     estimator = XGBRegressor(
        #         learning_rate=0.1, 
        #         n_estimators=50, 
        #         silent=False, 
        #         objective='reg:gamma',
        #         metric={'mae'}
        #         ),
        #     scoring = 'neg_mean_absolute_error',
        #     param_grid=param_grid,
        #     cv=3
        #     )
        # gsearch.fit(train_data, train_labels)
        # self.model = gsearch.best_estimator_
        self.model = XGBRegressor(n_estimators = 50,max_depth=50)
        self.model.fit(train_data, train_labels)
    
        
    def estimate_cor(self,test_labels,predict_test):
        cor = np.corrcoef(test_labels,predict_test)
        return cor
    
    def save_model(self,model_save_path):
        # effect = "Xgboost Mean Abs Error: %f " % (abs(gsearch.best_score_))
        joblib.dump(self.model, model_save_path, compress = 1) # Only best parameters

       
    
    def load_model(self,model_save_path):
        self.model = joblib.load(model_save_path)
        
        
    def predict_data(self,input_data):
        pred_data = self.model.predict(input_data)
        pred_data = pred_data.reshape(len(input_data))
        return pred_data

if __name__ =="__main__":
    model = xgb_gbdt()
    train_data = np.array([[1,2,3],[4,5,6],[7,8,9],
                           [1,2,3],[4,5,6],[7,8,9],
                           [1,2,3],[4,5,6],[7,8,9]])
    train_labels =np.array([0,1,0,0,1,0,0,1,0])
    
    model.train(train_data,train_labels)
    prediction_labels = model.predict_data(train_data)
    cor = model.estimate_cor(train_labels, prediction_labels)



